import pandas as pd
import time
import numpy as py
pd.set_option('display.max_rows', None)
start_time = time.time()
names = ['userId', 'movieId', 'rating', 'timestamp']
data = pd.read_csv('u.data', sep='\t', names=names)
data = data.sort_values(['movieId', 'userId'], ascending=True)
print(data.head())
n_users = data.userId.unique().shape[0]
n_items = data.movieId.unique().shape[0]
print(n_users, n_items)
data = data.drop(['timestamp'], axis=1)

# 获取数据的行数
data_rows = data.shape[0]
print('数据行数：', data.shape[0])
user_list_i = []
user_list_f = []
# 判断两个item是有共同的用户，并统计有多少个共同的用户
# print(data.loc[data['movieId'].isin([1, 3])].groupby('userId').size().isin([2]).sum(), 666)
# data_num = data.loc[data['movieId'].isin([1, 3])].groupby(['userId']).size().reset_index(name='Size')
# print(data_num.loc[data_num['Size'].isin([2])])
# print(data_num.loc[data_num['Size'].isin([2])].shape[0])

# *********************************************************
# # 统计用户1和用户2都看过的电影的部数
# data_num = data.loc[data['userId'].isin([1, 2])].groupby(['movieId']).size().reset_index(name='Size')
# print(data_num.loc[data_num['Size'].isin([2])])
# # 排序
# data_num = data_num.sort_values(ascending=False)
# print(data_num)

# 判断userId列里是否包含1和2，然后将符合条件的数据提取出来
# print(data.loc[data['userId'].isin([1, 2])])
# *********************************************************
# 计算两个用户之间的皮尔森相关系数
def pearsonr(movie_j, movie_q):
    """
    :param movie_j: 电影j
    :param movie_q: 电影q
    :return: 电影j和电影q的皮尔森相关系数
    """
    dif_sum = 0
    i_sum = 0
    f_sum = 0
    if data.loc[data['movieId'].isin([movie_j, movie_q])].groupby('userId').size().isin([2]).sum() < 8:
        return 0
    else:
        data_num = data.loc[data['movieId'].isin([movie_j, movie_q])].groupby(['userId']).size().reset_index(name='Size')
        data_num = data_num.loc[data_num['Size'].isin([2])]
        data_length = data_num.shape[0]
        ave_moviej_rating = data.loc[data['movieId'].isin([movie_j])]
        ave_moviej_rating = ave_moviej_rating['rating'].mean()
        ave_movieq_rating = data.loc[data['movieId'].isin([movie_j])]
        ave_movieq_rating = ave_movieq_rating['rating'].mean()

        for i in range(0, data_length):
            movie_j_user = data.loc[data['movieId'].isin([movie_j]) & data['userId'].isin([data_num.iloc[i][0]])]
            movie_q_user = data.loc[data['movieId'].isin([movie_q]) & data['userId'].isin([data_num.iloc[i][0]])]
            movie_j_rating = movie_j_user.iloc[0][2]
            movie_q_rating = movie_q_user.iloc[0][2]
            rating_i = movie_j_rating - ave_moviej_rating
            rating_f = movie_q_rating - ave_movieq_rating
            dif_sum += round(rating_i, 2) * round(rating_f, 2)
            i_sum += round(rating_i, 2) ** 2
            f_sum += round(rating_f, 2) ** 2
        if i_sum ** 0.5 * f_sum ** 0.5 == 0:
            pearson = 0
        else:
            pearson = dif_sum / (i_sum ** 0.5 * f_sum ** 0.5)
            pearson = (pearson + 1) / 2
        return pearson

# 电影i和其他电影的相似度
df = pd.DataFrame()
for i in range(1, n_items + 1):
    sim_dict = {}
    for j in range(1, n_items + 1):
        sim = pearsonr(i, j)
        if sim != 0:
            sim_dict[j] = sim
    if len(sim_dict) < 10:
        print(i, '相似电影少于10个')
        pass
    else:
        dict2 = sorted(sim_dict.items(), key=lambda x: x[1])
        print(dict2)
        sim_data = dict2[0:5] + dict2[-5:]
        df['movie%d' % i] = sim_data
        print(i)
# df.to_csv(path_or_buf="moviePear%d.csv" % i, index=False)
print(df)

end_time = time.time()
print('运行时间：', end_time - start_time, '秒')


